package demo

import (
	"bytes"
	"encoding/binary"
	"io"
	"log"
	"os"

	"gitee.com/general252/ebml-go/core"
	"gitee.com/general252/ebml-go/core/codecs"
	"gitee.com/general252/ebml-go/core/formats"
	"gitee.com/general252/ebml-go/mkv"
	"gitee.com/general252/ebml-go/mkvcore"
	"gitee.com/general252/ebml-go/webm"
	"github.com/bluenviron/mediacommon/pkg/codecs/h264"
	"github.com/bluenviron/mediacommon/pkg/codecs/mpeg4audio"
)

func H264AACDemo() {
	var inFilenameVideo = "input.h264"
	var inFilenameAudio = "output.aac"
	var outFilename = "outH264AAC.mkv"

	w, err := os.OpenFile(outFilename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666)
	if err != nil {
		log.Println(err)
		return
	} else {
		defer w.Close()
	}

	const useStartCode = true

	ws, err := webm.NewSimpleBlockWriter(w,
		[]*webm.TrackEntry{
			{
				Name:            "Video",
				TrackNumber:     1,
				TrackUID:        1111,
				CodecID:         core.VideoCodecMPEG4ISOAVC,
				TrackType:       1,
				DefaultDuration: 40000000,
				CodecPrivate:    nil,
				Video: &webm.Video{
					PixelWidth:  640,
					PixelHeight: 480,
				},
			},
			{
				Name:         "Audio",
				TrackNumber:  2,
				TrackUID:     2222,
				CodecID:      core.AudioCodecAAC,
				TrackType:    2,
				CodecPrivate: nil,
				Audio: &webm.Audio{
					SamplingFrequency: 8000,
					Channels:          1,
				},
			},
		},
		mkvcore.WithSeekHead(true),
		mkvcore.WithCues(true),
		mkvcore.WithEBMLHeader(mkv.DefaultEBMLHeader),
		mkvcore.WithSegmentInfo(mkv.DefaultSegmentInfo),
		mkvcore.WithMaxKeyframeInterval(1, 900*0x6FFF),
		mkvcore.WithSimpleTags(
			mkvcore.Tag{
				Targets: mkvcore.Targets{
					TagTrackUID: []uint64{
						1111,
					},
				},
				SimpleTag: []mkvcore.SimpleTag{
					{
						TagName:   "company",
						TagString: "xxx",
					},
				},
			},
			mkvcore.Tag{
				Targets: mkvcore.Targets{},
				SimpleTag: []mkvcore.SimpleTag{
					{
						TagName:   "xyz",
						TagString: "12585644567823132484",
					},
				},
			},
		),
	)
	if err != nil {
		log.Println(err)
		return
	} else {
		defer ws[0].Close()
		defer ws[1].Close()
	}

	videoStream := ws[0]
	audioStream := ws[1]

	/////////////////////////
	var fpVideo io.ReadCloser
	if fpVideo, err = os.Open(inFilenameVideo); err != nil {
		log.Println(err)
		return
	} else {
		defer fpVideo.Close()
	}

	var fpAudio io.ReadCloser
	if fpAudio, err = os.Open(inFilenameAudio); err != nil {
		log.Println(err)
		return
	} else {
		defer fpAudio.Close()
	}

	var (
		adtsList [][]byte
		index    int
	)
	formats.EmitADTSReader(fpAudio, func(packet *mpeg4audio.ADTSPacket, data, _ []byte) {
		adtsList = append(adtsList, data)
	})
	var getAudioData = func(index int) ([]byte, bool) {
		if index >= 0 && index < len(adtsList) {
			return adtsList[index], true
		}
		return nil, false
	}

	var (
		timestamp int64
		startCode = []byte{0x0, 0x0, 0x0, 0x01}
		ss        []h264.NALUType
		sps       = &bytes.Buffer{}
		pps       = &bytes.Buffer{}
	)

	var getBlockData = func(t h264.NALUType, data []byte) *bytes.Buffer {
		switch t {
		case h264.NALUTypeSPS:
			sps.Reset()
			if useStartCode {
				sps.Write(startCode)
			} else {
				_ = binary.Write(sps, binary.BigEndian, int32(len(data)))
			}
			sps.Write(data)
			return nil
		case h264.NALUTypePPS:
			pps.Reset()
			if useStartCode {
				sps.Write(startCode)
			} else {
				_ = binary.Write(sps, binary.BigEndian, int32(len(data)))
			}
			pps.Write(data)
			return nil
		case h264.NALUTypeIDR:
			var buff = bytes.Buffer{}
			buff.Write(sps.Bytes())
			buff.Write(pps.Bytes())

			if useStartCode {
				buff.Write(startCode)
			} else {
				_ = binary.Write(&buff, binary.BigEndian, int32(len(data)))
			}
			buff.Write(data)
			return &buff
		case h264.NALUTypeNonIDR:
			var buff = bytes.Buffer{}
			if useStartCode {
				buff.Write(startCode)
			} else {
				_ = binary.Write(&buff, binary.BigEndian, int32(len(data)))
			}
			buff.Write(data)
			return &buff
		}

		return nil
	}

	// r, err := h264reader.NewReader(fp)
	codecs.EmitNALUH264Reader(fpVideo, codecs.NALUFormatAnnexB, codecs.NALUFormatNo, func(t h264.NALUType, data []byte) {
		ss = append(ss, t)

		switch t {
		case h264.NALUTypeNonIDR: // Coded slice of a non-IDR picture
			// 1
			buff := getBlockData(t, data)
			_, _ = videoStream.Write(false, timestamp, buff.Bytes())

			if audioData, ok := getAudioData(index); ok {
				index++
				_, _ = audioStream.Write(true, timestamp, audioData)
			}
			timestamp += 40
		case h264.NALUTypeDataPartitionA: // Coded slice data partition A
		case h264.NALUTypeDataPartitionB: // Coded slice data partition B
		case h264.NALUTypeDataPartitionC: // Coded slice data partition C
		case h264.NALUTypeIDR: // Coded slice of an IDR picture
			// 5
			buff := getBlockData(t, data)
			_, _ = videoStream.Write(true, timestamp, buff.Bytes())

			if audioData, ok := getAudioData(index); ok {
				index++
				_, _ = audioStream.Write(true, timestamp, audioData)
			}
			timestamp += 40
		case h264.NALUTypeSEI: // Supplemental enhancement information (SEI)
			// 6
		case h264.NALUTypeSPS: // Sequence parameter set
			// 7
			getBlockData(t, data)
			// _, _ = videoStream.Write(false, timestamp, append(startCode, bytes...))
		case h264.NALUTypePPS: // Picture parameter set
			// 8
			getBlockData(t, data)
		case h264.NALUTypeAccessUnitDelimiter: // Access unit delimiter
		case h264.NALUTypeEndOfSequence: // End of sequence
		case h264.NALUTypeEndOfStream: // End of stream
		case h264.NALUTypeFillerData: // Filler data
		case h264.NALUTypeSPSExtension: // Sequence parameter set extension
		case h264.NALUTypeSliceLayerWithoutPartitioning: // Coded slice of an auxiliary coded picture without partitioning:

		}
	})

	log.Println(ss)
}
